package com.demo;

import java.io.*;
import java.net.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;
import com.demo.node.NodeType;

/**
 * P2P网络管理类
 * 负责处理节点间的通信、交易广播和区块同步
 */
public class P2PNetwork {
    private int port;
    private ServerSocket serverSocket; // 用于接收其他节点的连接请求
    private Chain blockchain;
    private final Map<String, ClientHandler> peerHandlers = new ConcurrentHashMap<>(); // 使用线程安全的 ConcurrentHashMap 存储所有对等节点的处理器  // 键为 "ip:port", 值为该连接的处理器
    private ExecutorService threadPool; // 线程池，用于并发处理网络请求
    private NodeType nodeType; // 节点类型字段
    private String nodePublicKey; // 添加节点公钥字段

    // 修改构造函数，添加公钥参数
    public P2PNetwork(int port, Chain blockchain, NodeType nodeType, String nodePublicKey) {
        this.port = port;
        this.blockchain = blockchain;
        this.nodeType = nodeType;
        this.nodePublicKey = nodePublicKey;
        this.threadPool = Executors.newFixedThreadPool(20); // 增加线程池大小以处理更多持久连接
        startServer();
    }

    // 启动一个 ServerSocket，监听指定端口，并接受其他节点的连接请求。
    private void startServer() {
        threadPool.submit(() -> {
            try {
                serverSocket = new ServerSocket(port);
                String localIP = InetAddress.getLocalHost().getHostAddress();
                System.out.println("P2P节点启动成功！ IP: " + localIP + ", Port: " + port);
                while (!serverSocket.isClosed()) {
                    Socket clientSocket = serverSocket.accept();
                    // 为每一个新的连接创建一个处理器并由线程池管理
                    ClientHandler handler = new ClientHandler(clientSocket, this);
                    threadPool.submit(handler);
                }
            } catch (IOException e) {
                if (!serverSocket.isClosed()) {
                     System.out.println("服务器套接字出错: " + e.getMessage());
                }
            }
        });
    }

    // 主动连接另一个 P2P 节点
    public void connectToPeer(String ip, int peerPort) {
        String peerId = ip + ":" + peerPort;
        if (peerHandlers.containsKey(peerId)) {
            System.out.println("已存在到节点 " + peerId + " 的连接。");
            return;
        }
        try {
            Socket socket = new Socket(ip, peerPort);
            // 连接成功后，创建处理器并由线程池管理
            ClientHandler handler = new ClientHandler(socket, this);
            threadPool.submit(handler);
            // 发送握手消息，包含自己的监听端口
            handler.send("HANDSHAKE:" + this.port);
            // 连接成功后，请求区块链同步
            syncBlockchain(handler);
        } catch (IOException e) {
            System.out.println("连接节点失败: " + ip + ":" + peerPort + " - " + e.getMessage());
            e.printStackTrace();
        }
    }

    // 向所有对等节点广播消息，可选择排除一个源头处理器以避免消息回传
    public void broadcast(Object message, ClientHandler excludeHandler) {
        for (ClientHandler handler : peerHandlers.values()) {
            if (handler != excludeHandler) {
                handler.send(message);
            }
        }
    }
    
    // 对外广播接口
    public void broadcast(Object message) {
        broadcast(message, null);
    }

    // ClientHandler在握手成功后调用此方法进行注册
    public void registerPeer(String peerId, ClientHandler handler) {
        peerHandlers.put(peerId, handler);
        System.out.println("新节点已连接: " + peerId + "。当前总连接数: " + peerHandlers.size());
    }

    // ClientHandler在连接断开时调用此方法进行注销
    public void unregisterPeer(String peerId) {
        if (peerId != null) {
            peerHandlers.remove(peerId);
            System.out.println("节点已断开: " + peerId + "。当前总连接数: " + peerHandlers.size());
        }
    }

    /**
     * @Description: 处理单个对等节点连接的处理器，负责所有I/O操作
     * @Version: 2.0
     **/
    private class ClientHandler implements Runnable {
        private Socket socket;
        private ObjectInputStream in;
        private ObjectOutputStream out;
        private String peerId;
        private P2PNetwork network;

        public ClientHandler(Socket socket, P2PNetwork network) {
            this.socket = socket;
            this.network = network;
            try {
                this.out = new ObjectOutputStream(socket.getOutputStream());
                this.in = new ObjectInputStream(socket.getInputStream());
            } catch (IOException e) {
                System.out.println("创建处理器时发生异常: " + e.getMessage());
            }
        }

        public void send(Object message) {
            try {
                out.writeObject(message);
                out.flush();
            } catch (IOException e) {
                System.out.println("发送消息时发生异常: " + e.getMessage());
            }
        }

        @Override
        public void run() {
            try {
                while (!socket.isClosed()) {
                    Object message = in.readObject();
                    if (message instanceof String) {
                        String msg = (String) message;
                        if (msg.startsWith("HANDSHAKE:")) {
                            // 处理握手消息
                            int remotePort = Integer.parseInt(msg.split(":")[1]);
                            String remoteIP = socket.getInetAddress().getHostAddress();
                            peerId = remoteIP + ":" + remotePort;
                            network.registerPeer(peerId, this);
                        } else if (msg.startsWith("SYNC_REQUEST:")) {
                            // 解析请求方节点类型
                            NodeType requesterType = NodeType.valueOf(msg.split(":")[1]);
                            network.handleSyncRequest(this, requesterType);
                        } else if (msg.equals("SYNC_RESPONSE")) {
                            // 读取完整的区块链数据
                            List<Block> receivedChain = (List<Block>) in.readObject();
                            network.handleSyncResponse(receivedChain, this);
                        } else if (msg.equals("HEADER_SYNC_RESPONSE")) {
                            // 读取区块头信息
                            List<Block> headers = (List<Block>) in.readObject();
                            network.handleSyncResponse(headers, this);
                        }
                    } else if (message instanceof Transaction) {
                        Transaction tx = (Transaction) message;
                        try {
                            // 使用节点的公钥而不是从blockchain获取
                            String myPubKey = network.getNodePublicKey();
                            if (myPubKey != null && tx.getTo().equals(myPubKey)) {
                                network.blockchain.addTxToPool(tx);
                                System.out.println("收到新的交易，金额: " + tx.amount + " BTC，等待验证...");
                            }
                            // 广播交易给其他节点
                            network.broadcast(tx, this);
                        } catch (Exception e) {
                            System.out.println("处理交易时发生异常: " + e.getMessage());
                            e.printStackTrace();
                        }
                    } else if (message instanceof Block) {
                        Block block = (Block) message;
                        try {
                            // 验证区块中的所有交易
                            boolean isValid = true;
                            for (Transaction tx : block.transactions) {
                                if (!tx.validate()) {
                                    isValid = false;
                                    break;
                                }
                            }
                            // 如果所有交易都验证通过，且链验证通过，则添加区块
                            if (isValid && network.blockchain.validateChain()) {
                                network.blockchain.getChain().add(block);
                                network.broadcast(block, this);
                            }
                        } catch (Exception e) {
                            System.out.println("处理区块时发生异常: " + e.getMessage());
                        }
                    }
                }
            } catch (EOFException e) {
                System.out.println("连接已关闭: " + (peerId != null ? peerId : "未知节点"));
            } catch (Exception e) {
                System.out.println("处理消息时发生异常: " + e.getMessage());
                e.printStackTrace();
            } finally {
                closeConnection();
            }
        }

        public void closeConnection() {
            try {
                if (peerId != null) {
                    network.unregisterPeer(peerId);
                }
                if (in != null) in.close();
                if (out != null) out.close();
                if (socket != null) socket.close();
            } catch (IOException e) {
                System.out.println("关闭连接时发生异常: " + e.getMessage());
            }
        }
    }

    // 验证交易后广播到整个网络。
    public void broadcastTransaction(Transaction tx) {
        System.out.println("广播交易");
        try {
            // 支付方发起交易时不需要验证，直接广播
            blockchain.addTxToPool(tx);
            broadcast(tx);
        } catch (Exception e) {
            System.out.println("广播交易时发生异常: " + e.getMessage());
            e.printStackTrace();
        }
    }

    // 验证区块后广播到整个网络。
    public void broadcastBlock(Block block) {
        System.out.println("广播区块");
        try {
            broadcast(block);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // 优雅关闭线程池和网络连接。
    public void shutdown() {
        System.out.println("正在关闭P2P网络...");
        try {
            if (serverSocket != null) serverSocket.close();
            for (ClientHandler handler : peerHandlers.values()) {
                handler.closeConnection();
            }
            peerHandlers.clear();
            threadPool.shutdown();
            if (!threadPool.awaitTermination(5, TimeUnit.SECONDS)) {
                threadPool.shutdownNow();
            }
        } catch (InterruptedException | IOException e) {
            threadPool.shutdownNow();
        }
        System.out.println("P2P网络已关闭。");
    }

    // 在握手成功后进行数据同步
    public void syncBlockchain(ClientHandler handler) {
        System.out.println("\n=== 请求区块链同步 ===");
        System.out.println("当前节点类型: " + nodeType);
        System.out.println("当前区块链长度: " + blockchain.getChain().size());
        
        // 发送同步请求，包含节点类型信息
        handler.send("SYNC_REQUEST:" + nodeType);
        System.out.println("===================\n");
    }

    // 处理同步请求
    private void handleSyncRequest(ClientHandler handler, NodeType requesterType) {
        System.out.println("\n=== 处理同步请求 ===");
        System.out.println("请求方节点类型: " + requesterType);
        System.out.println("当前节点类型: " + nodeType);
        
        try {
            if (requesterType == NodeType.LIGHT_NODE) {
                // 只发送区块头信息
                System.out.println("发送区块头信息");
                System.out.println("当前区块链高度: " + blockchain.getChain().size());
                List<Block> headers = new ArrayList<>();
                for (Block block : blockchain.getChain()) {
                    // 创建一个只包含必要信息的区块头
                    Block header = new Block(new ArrayList<>(), block.prevHash);
                    header.hash = block.hash;
                    header.time = block.time;
                    header.nonce = block.nonce;
                    headers.add(header);
                }
                handler.send("HEADER_SYNC_RESPONSE");
                handler.send(headers);
            } else {
                // 发送完整区块链数据
                System.out.println("发送完整区块链数据");
                System.out.println("当前区块链高度: " + blockchain.getChain().size());
                handler.send("SYNC_RESPONSE");
                handler.send(new ArrayList<>(blockchain.getChain()));
            }
            System.out.println("同步数据发送完成");
            System.out.println("===================\n");
        } catch (Exception e) {
            System.out.println("处理同步请求时发生异常: " + e.getMessage());
            e.printStackTrace();
        }
    }

    // 处理同步响应
    private void handleSyncResponse(List<Block> receivedChain, ClientHandler handler) {
        if (receivedChain != null && !receivedChain.isEmpty()) {
            try {
                System.out.println("\n=== 处理同步响应 ===");
                System.out.println("收到区块链数据，高度: " + receivedChain.size());
                
                // 验证接收到的区块链
                if (blockchain.validateChain()) {
                    // 更新本地区块链
                    blockchain.getChain().clear();
                    blockchain.getChain().addAll(receivedChain);
                    System.out.println("区块链同步成功，当前高度: " + blockchain.getChain().size());
                    // 打印同步后的区块链信息
                    System.out.println("\n=== 同步后的区块链信息 ===");
                    System.out.println(blockchain.toString());
                    System.out.println("========================\n");
                } else {
                    System.out.println("区块链验证失败，同步中止");
                }
            } catch (Exception e) {
                System.out.println("区块链同步验证失败: " + e.getMessage());
                e.printStackTrace();
            }
        } else {
            System.out.println("收到的区块链数据为空");
        }
    }

    public NodeType getNodeType() {
        return nodeType;
    }

    // 添加获取节点公钥的方法
    public String getNodePublicKey() {
        return nodePublicKey;
    }
}